import re
from .CATAGORY import *
import os

# 获取当前脚本所在的目录路径
script_dir = os.path.dirname(os.path.abspath(__file__))
absolute_path = ('敏感词汇表.txt')

# 获取相对路径
file_path = os.path.join(script_dir,absolute_path)


class Promptlist_check(Ac):
    mode = ['auto','custom']
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
            "mode":(s.mode,),
            "text": ("STRING",{"forceInput":True}),
            "sensitive_words_list": ("STRING", {
                "multiline": True, 
                "default": ''}),
            
        }}
    # 返回结果类型
    RETURN_TYPES = ('STRING',)
    
    # 返回节点命名
    RETURN_NAMES = ('STRING',)
    FUNCTION = "promptlist_check" 

    @classmethod
    def read_and_split_txt(self,file_path):
        with open(file_path, 'r') as file:
            content = file.read()
            split_content = content.split(',')
            return [item.strip() for item in split_content]

    def promptlist_check(self,mode,text,sensitive_words_list):
        if mode == 'auto':
            sensitive_words_list = self.read_and_split_txt(file_path)
        if mode == 'custom':
            sensitive_words_list = sensitive_words_list.split(",")
        for word in sensitive_words_list:
            pattern = re.compile(r'\b' + re.escape(word) + r'\b')
            text = re.sub(pattern, '', text)
        result = text.strip()
        return (result,)


class Promptlist_check_rp(Ac):
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
            "input_string": ("STRING",{"forceInput":True}),
            "filter_word": ("STRING", {
                "multiline": True, 
                "default": '这里主要是输入需要被过滤或者替换的词语'}),
            
        }}
    # 返回结果类型
    RETURN_TYPES = ('STRING',)
    
    # 返回节点命名
    RETURN_NAMES = ('STRING',)
    FUNCTION = "promptlist_check" 


    def promptlist_check(self,input_string,filter_word):
        result = input_string.replace(filter_word, "")
        return (result,)        


# ==============================================================
if __name__ == "__main__":

    pass

